# main.py
from importlib.util import LazyLoader
from experiment_recorder import ExperimentRecorder
from train_unet0 import UNetTrainer
import torch

def main():
    # 初始化实验记录器
    recorder = ExperimentRecorder(
        experiment_name="训练加，看看中心损失能否降低",
        base_dir="/data/zhouhai/fuxian/实验记录"
    )
    
    # 设置设备
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 记录实验参数
    recorder.log_parameters({
        'batch_size': 5,
        'epochs': 20,
        'learning_rate': 0.01,
        'target_point': (230,135),
        'center_point': (256, 256),
        'max_layers': 9,
        'device': str(device)
    })
    
    # 初始化训练器
    trainer = UNetTrainer(recorder, device)
    
    # 设置数据
    print("\n正在加载数据集...")
    trainer.setup_data(
        root_path="/data/zhouhai/fuxian/视/处理数据集/图像定位",
        batch_size=5
    )
    
    # 训练各层模型
    print("\n开始训练...")
    trained_models = []
    for layer in range(0,10):  # 0到9层
        print(f"\n训练层 {layer}...")
        model = trainer.train_layer(
            layer=layer,
            epochs=20,
            target_point=(230, 135),
            center_point=(256, 256),
            max_layers=9
        )
        trained_models.append(model)
    
    # 评估模型
    print("\n评估模型...")
    results_df = trainer.evaluate(trained_models, target_point=(230, 135))
    
    # 完成实验记录
    recorder.finalize()
    
    print("\n实验全部完成！")

if __name__ == "__main__":
    main()